import subprocess


# 通过ssh连接并使用rsync上传文件
def upload_file_via_rsync(local_path, remote_path, server_info):
    command = [
        'rsync', '-avz', local_path,
        f"{server_info['username']}@{server_info['host']}:{remote_path}"
    ]
    subprocess.run(command)


# 通过ssh连接并使用rsync下载文件
def download_file_via_rsync(remote_path, local_path, server_info):
    command = [
        'rsync', '-avz',
        f"{server_info['username']}@{server_info['host']}:{remote_path}",
        local_path
    ]
    subprocess.run(command)


# 选择要连接的服务器并返回服务器信息
def select_server():
    # 配置太空服务器A机（训练卡）和B机（推理卡）
    servers = {
        "A": {
            "host": "192.168.200.103",  # 太空服务器A机IP
            "username": "user",
            "password": "user"
        },
        "B": {
            "host": "192.168.200.108",  # 太空服务器B机IP
            "username": "user",
            "password": "user"
        }
    }

    selected_server = input("选择要连接的服务器 (A: 训练卡, B: 推理卡): ").upper()

    if selected_server in servers:
        return servers[selected_server]
    else:
        print("无效选择，请选择 A 或 B")
        return None


# 主程序
if __name__ == '__main__':
    # 配置跳板树莓派的连接信息
    jump_host = "192.168.100.165"  # 树莓派跳板IP
    jump_username = "ubuntu"  # 树莓派跳板用户名
    jump_password = "123456"  # 树莓派跳板密码

    # 选择要连接的服务器
    server_info = select_server()

    if server_info:
        # 上传船的照片
        local_photo_path = "/path/to/boat/photo.jpg"  # 本地船的照片路径
        remote_photo_path = "/home/user/experiment/16/work/image"  # 远程船的照片上传路径
        upload_file_via_rsync(local_photo_path, remote_photo_path, server_info)
        print(f"船的照片上传到服务器 {server_info['host']} 完成")

        # 上传 ResNet 模型
        local_model_path = "/path/to/resnet/model.pth"  # 本地 ResNet 模型文件路径
        remote_model_path = "/home/user/experiment/16/work"  # 远程 ResNet 模型上传路径
        upload_file_via_rsync(local_model_path, remote_model_path, server_info)
        print(f"ResNet 模型上传到服务器 {server_info['host']} 完成")

        # 下载文件
        remote_download_file = "/path/on/server/file"  # 远程文件路径
        local_download_path = "/path/to/local/download"  # 本地下载路径
        download_file_via_rsync(remote_download_file, local_download_path, server_info)
        print(f"从服务器 {server_info['host']} 下载文件完成")
